﻿import torch
import torch.nn as nn
import matplotlib.pyplot as plt
from Dataset.MyDataset.Regression.linear_regression_dataset import (
    get_linear_data,
)

X, y = get_linear_data(1)


class MyLinearModel(nn.Module):
    def __init__(self, in_fea, out_fea):
        super(MyLinearModel, self).__init__()
        self.out = nn.Linear(in_features=in_fea, out_features=out_fea)

    def forward(self, x):
        x = self.out(x)
        return x


input_x = torch.tensor(X).to(torch.float32)
input_y = torch.unsqueeze(torch.from_numpy(y).float(), dim=1)

model = MyLinearModel(1, 1)
loss_func = nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.02)

plt.ion()
for step in range(9000):
    pred = model(input_x)
    loss = loss_func(pred, input_y)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    if step % 10 == 0:
        plt.cla()
        plt.scatter(X, y) # 样本点
        plt.plot(
            input_x.squeeze().data.numpy(), pred.squeeze().data.numpy(), "r-", lw=5
        ) # 拟合直线
        [w, b] = model.parameters()
        plt.text(0, 0, "loss=%.4f, k=%.2f, b=%.2f" % (loss.item(), w.item(), b.item()))
        plt.pause(0.1)
plt.ioff()
plt.show()